// Copyright 2022 Chen Jun
// Licensed under the MIT License.

#ifndef ARMOR_DETECTOR_PNP_SOLVER_HPP_
#define ARMOR_DETECTOR_PNP_SOLVER_HPP_

#include "hnurm_armor/Compensator.h"
#include "hnurm_armor/DataType.hpp"
#include "hnurm_armor/armor.hpp"
#include "hnurm_armor/tracker.hpp"

#include <opencv2/calib3d.hpp>
#include <opencv2/core.hpp>
#include <opencv2/core/eigen.hpp>
#include <Eigen/Core>
#include <Eigen/Dense>

#include <angles/angles.h>

#include <cmath>
#include <cstdlib>
#include <geometry_msgs/msg/transform_stamped.hpp>
#include <tf2/LinearMath/Quaternion.h>
#include <tf2_eigen/tf2_eigen.hpp>
#include <vector>

namespace hnurm
{
    class TSolver
    {
    public:
        explicit TSolver(const rclcpp::Node::SharedPtr &node)
        {
            small_armor_points_ = {
                cv::Point3f(0, 0, SMALL_ARMOR_HEIGHT / 2),
                cv::Point3f(-SMALL_ARMOR_WIDTH / 2, 0, 0),
                cv::Point3f(0, 0, -SMALL_ARMOR_HEIGHT / 2),
                cv::Point3f(SMALL_ARMOR_WIDTH / 2, 0, 0),
                cv::Point3f(0, 0, 0),
            };
            small_armor_9points_ = {
                cv::Point3f(-SMALL_ARMOR_WIDTH / 2, 0, SMALL_ARMOR_HEIGHT / 2),
                cv::Point3f(-SMALL_ARMOR_WIDTH / 2, 0, -SMALL_ARMOR_HEIGHT / 2),
                cv::Point3f(SMALL_ARMOR_WIDTH / 2, 0, -SMALL_ARMOR_HEIGHT / 2),
                cv::Point3f(SMALL_ARMOR_WIDTH / 2, 0, SMALL_ARMOR_HEIGHT / 2),
                cv::Point3f(0, 0, SMALL_ARMOR_HEIGHT / 2),
                cv::Point3f(-SMALL_ARMOR_WIDTH / 2, 0, 0),
                cv::Point3f(0, 0, -SMALL_ARMOR_HEIGHT / 2),
                cv::Point3f(SMALL_ARMOR_WIDTH / 2, 0, 0),
                cv::Point3f(0, 0, 0),
            };
            large_armor_points_ = {
                cv::Point3f(0, 0, LARGE_ARMOR_WIDTH / 2),
                cv::Point3f(-LARGE_ARMOR_WIDTH / 2, 0, 0),
                cv::Point3f(0, 0, -LARGE_ARMOR_WIDTH / 2),
                cv::Point3f(LARGE_ARMOR_WIDTH / 2, 0, 0),
                cv::Point3f(0, 0, 0),
            };
            large_armor_9points_ = {
                cv::Point3f(-LARGE_ARMOR_WIDTH / 2, 0, LARGE_ARMOR_HEIGHT / 2),
                cv::Point3f(-LARGE_ARMOR_WIDTH / 2, 0, -LARGE_ARMOR_HEIGHT / 2),
                cv::Point3f(LARGE_ARMOR_WIDTH / 2, 0, -LARGE_ARMOR_HEIGHT / 2),
                cv::Point3f(LARGE_ARMOR_WIDTH / 2, 0, LARGE_ARMOR_HEIGHT / 2),
                cv::Point3f(0, 0, LARGE_ARMOR_WIDTH / 2),
                cv::Point3f(-LARGE_ARMOR_WIDTH / 2, 0, 0),
                cv::Point3f(0, 0, -LARGE_ARMOR_WIDTH / 2),
                cv::Point3f(LARGE_ARMOR_WIDTH / 2, 0, 0),
                cv::Point3f(0, 0, 0),
            };
            float cam_bias_z, cam_bias_y;
            cam_bias_z = node->declare_parameter("processor.solver.camera_bias_z", 0);
            cam_bias_y = node->declare_parameter("processor.solver.camera_bias_y", 0);
            _cam_bias << cam_bias_z, cam_bias_y, 0;
            std::string calib_info_path;
            // cfg_node["calib_info_path"] >> calib_info_path;
            calib_info_path = node->declare_parameter("processor.solver.calib_info_path", "src/hnurm_armor/data/");
            std::string camera_id;
            camera_id = node->declare_parameter("processor.solver.camera_id", "0");
            calib_info_path = calib_info_path + camera_id + ".yaml";
            cv::FileStorage calib_info(calib_info_path, cv::FileStorage::READ);
            if (!calib_info.isOpened())
            {
                RCLCPP_WARN_STREAM(node->get_logger(), "Failed to load calib info in " << calib_info_path);
            }
            calib_info["camera_matrix"] >> camera_matrix_;
            calib_info["distortion_coefficients"] >> dist_coeffs_;
        }
        float distance(const cv::Point2f &p1, const cv::Point2f &p2)
        {
            float dx = p1.x - p2.x;
            float dy = p1.y - p2.y;
            return sqrt(dx * dx + dy * dy);
        }

        float totalDistance(const std::vector<cv::Point2f> &points1, const std::vector<cv::Point2f> &points2)
        {
            float total_dist = 0.0f;
            for (size_t i = 0; i < points1.size(); ++i)
            {
                total_dist += distance(points1[i], points2[i]);
            }
            return total_dist;
        }
        static void rotationMatrixToEulerAngles(const Eigen::Matrix3d &R, Eigen::Vector3d &eulerAngles)
        {
            double sy = sqrt(R(0, 0) * R(0, 0) + R(1, 0) * R(1, 0));
            bool singular = sy < 1e-6; // If
            if (!singular)
            {
                eulerAngles[0] = atan2(R(2, 1), R(2, 2)); // pitch
                eulerAngles[1] = atan2(-R(2, 0), sy); // roll
                eulerAngles[2] = atan2(R(1, 0), R(0, 0)); // yaw
            }
            else
            {
                eulerAngles[0] = atan2(-R(1, 2), R(1, 1));
                eulerAngles[1] = atan2(-R(2, 0), sy);
                eulerAngles[2] = 0;
            }
        }

        float get_distance(const float &yaw_reg, Armor &armor, cv::Mat &tvec)
        {
            double pitch_deg;
            if (armor.number == "outpost")
            {
                pitch_deg = 15.0;
            }
            else if (armor.number == "base")
            {
                pitch_deg = -62.5;
            }
            else
            {
                pitch_deg = -15.0;
            }
            double pitch_rad = pitch_deg * CV_PI / 180.0;
            double yaw_rad = yaw_reg * CV_PI / 180.0;
            Eigen::Matrix3d rotation_matrix_pitch_roll;
            rotation_matrix_pitch_roll << 1, 0, 0,
                0, cos(pitch_rad), -sin(pitch_rad),
                0, sin(pitch_rad), cos(pitch_rad);
            Eigen::Matrix3d rotation_matrix_yaw;
            rotation_matrix_yaw << cos(yaw_rad), -sin(yaw_rad), 0,
                sin(yaw_rad), cos(yaw_rad), 0,
                0, 0, 1;
            Eigen::Matrix3d final_rotation_matrix = rotation_matrix_yaw * rotation_matrix_pitch_roll;
            std::vector<cv::Point2f> image_points;
            cv::Matx33d final_rotation_matrix_opencv(
                final_rotation_matrix(0, 0), final_rotation_matrix(0, 1), final_rotation_matrix(0, 2),
                final_rotation_matrix(1, 0), final_rotation_matrix(1, 1), final_rotation_matrix(1, 2),
                final_rotation_matrix(2, 0), final_rotation_matrix(2, 1), final_rotation_matrix(2, 2)
            );
            cv::Matx33d r;
            cv::eigen2cv(armor._R, r);
            if (armor.armor_type == ArmorType::SMALL)
            {
                projectPoints(small_armor_9points_,
                              r * final_rotation_matrix_opencv,
                              tvec,
                              camera_matrix_,
                              dist_coeffs_,
                              image_points);
            }
            else
            {
                projectPoints(large_armor_9points_,
                              r * final_rotation_matrix_opencv,
                              tvec,
                              camera_matrix_,
                              dist_coeffs_,
                              image_points);
            }
//        std::vector<cv::Point2f> last_five_points(armor.points2d.end() - 5, armor.points2d.end());
            float distance_9 = totalDistance(armor.points2d, image_points);
            return distance_9;
        }
        // Get 3d position from 2d points
        bool GetTransformation(Armor &armor, float g_pitch, float g_yaw)
        {
            // 求解PnP问题
            //            g_pitch = 0;
            //                    g_yaw=0;
            cv::Mat rvec, tvec;
            if (armor.armor_type == ArmorType::SMALL)
            {
                cv::solvePnP(
                    small_armor_9points_,
                    armor.points2d,
                    camera_matrix_,
                    dist_coeffs_,
                    rvec,
                    tvec,
                    false,
                    cv::SOLVEPNP_IPPE
                );
//            solvePnPBA(small_armor_points_, armor.points2d, camera_matrix_, dist_coeffs_, rvec, tvec, true);
            }
            else if (armor.armor_type == ArmorType::LARGE)
                cv::solvePnP(
                    large_armor_9points_,
                    armor.points2d,
                    camera_matrix_,
                    dist_coeffs_,
                    rvec,
                    tvec,
                    false,
                    cv::SOLVEPNP_IPPE
                );

            // 构造变换矩阵
            armor._translation << tvec.at<double>(0, 0), tvec.at<double>(1, 0), tvec.at<double>(2, 0);
            Eigen::AngleAxisd rot_angle(
                cv::norm(rvec), Eigen::Vector3d(rvec.at<double>(0, 0), rvec.at<double>(1, 0), rvec.at<double>(2, 0))
            );
            armor._rmat = rot_angle.matrix();
            float sp = std::sin(g_pitch), cp = std::cos(g_pitch), sy = std::sin(g_yaw), cy = std::cos(g_yaw);
            armor._r1 << 1, 0, 0,
                0, 0, 1,
                0, -1, 0;
            armor._r2 << 1, 0, 0,
                0, cp, -sp,
                0, sp, cp;
            armor._r3 << cy, -sy, 0,
                sy, cy, 0,
                0, 0, 1;
            armor._R = (armor._r3 * armor._r2 * armor._r1).inverse();
            armor.rotation = armor._r3 * armor._r2 * armor._r1 * armor._rmat;
            armor.position = armor._r3 * armor._r2 * (armor._r1 * armor._translation) / 1'000;
//            Eigen::Vector3d eulerAngles;
//            rotationMatrixToEulerAngles(armor.rotation, eulerAngles);
//            float a = eulerAngles[2] * 180 / CV_PI - 35;  // 搜索区间起点
//            float b = eulerAngles[2] * 180 / CV_PI + 35;   // 搜索区间终点
//            float tolerance = 0.1; // 精度阈值
//            const float phi = (1 + std::sqrt(5)) / 2; // 黄金分割比
//            float c = b - (b - a) / phi;
//            float d = a + (b - a) / phi;
//            while (std::fabs(b - a) > tolerance)
//            {
//                float s1 = get_distance(c, armor, tvec);
//                float s2 = get_distance(d, armor, tvec);
//                if (s1 < s2)
//                {
//                    b = d;
//                }
//                else
//                {
//                    a = c;
//                }
//                c = b - (b - a) / phi;
//                d = a + (b - a) / phi;
//            }
//            float right_yaw = (a + b) / 2 * CV_PI / 180;
//            Eigen::Matrix3d rotation_matrix_pitch_roll;
//            double pitch_deg;
//            if (armor.number == "outpost")
//            {
//                pitch_deg = 15.0;
//            }
//            else if (armor.number == "base")
//            {
//                pitch_deg = -62.5;
//            }
//            else
//            {
//                pitch_deg = -15.0;
//            }
//            rotation_matrix_pitch_roll << 1, 0, 0,
//                0, cos(pitch_deg * CV_PI / 180.0), -sin(pitch_deg * CV_PI / 180.0),
//                0, sin(pitch_deg * CV_PI / 180.0), cos(pitch_deg * CV_PI / 180.0);
//            Eigen::Matrix3d rotation_matrix_yaw;
//            rotation_matrix_yaw << cos(right_yaw), -sin(right_yaw), 0,
//                sin(right_yaw), cos(right_yaw), 0,
//                0, 0, 1;
//            armor.rotation = rotation_matrix_yaw * rotation_matrix_pitch_roll;
            return true;
        }


    private:
        cv::Mat camera_matrix_;
        cv::Mat dist_coeffs_;

        static constexpr float SMALL_ARMOR_WIDTH = 131.25;
        static constexpr float SMALL_ARMOR_HEIGHT = 55.5;
        static constexpr float LARGE_ARMOR_WIDTH = 226.25;
        static constexpr float LARGE_ARMOR_HEIGHT = 55.5;

        // Four vertices of armor in 3d
        std::vector<cv::Point3f> small_armor_points_;
        std::vector<cv::Point3f> small_armor_9points_;
        std::vector<cv::Point3f> large_armor_points_;
        std::vector<cv::Point3f> large_armor_9points_;
        Eigen::Vector3d _cam_bias;
        double angle = 0.2 * CV_PI;
        float lastr = 0.14;
        float radius = 0.24;
        cv::Point2f ccc = cv::Point2f(3, 3);
        std::vector<cv::Point3f> armorss;
        std::vector<cv::Point3f> armorss1;
    };

}  // namespace hnurm

#endif  // ARMOR_DETECTOR_PNP_SOLVER_HPP_
